#include "kpextComputePipeline.h"

using namespace kpext;


kpext::ComputePipeline::ComputePipeline(
        std::shared_ptr<vk::Device> device,
        std::shared_ptr<vk::PipelineCache> globaPipelinecache
        ) :
    m_descrSetLayouts({}),
    m_specConsts(nullptr),
    m_pushConsts(nullptr),
    m_useGlobalPipelinecache(true),
    m_spirvCode({}),
    m_shaderModule(nullptr),
    m_layout(nullptr),
    m_pipelineCache(nullptr),
    m_pipeline(nullptr),
    m_GPipelineCache(globaPipelinecache)
{
    if(!device){
        throw std::runtime_error(__FUNCTION__ " input null pointer!");
    }
    m_device = device;
}

ComputePipeline::~ComputePipeline(){
    if(!m_device) return;

    m_descrSetLayouts.clear();
    m_spirvCode.clear();

    if(m_pipeline){
        m_device->destroyPipeline(*m_pipeline);
        m_pipeline = nullptr;
    }

    if(m_pipelineCache){
        m_device->destroyPipelineCache(*m_pipelineCache);
        m_pipelineCache = nullptr;
    }

    if(m_layout){
        m_device->destroyPipelineLayout(*m_layout);
        m_layout = nullptr;
    }

    if(m_shaderModule){
        m_device->destroyShaderModule(*m_shaderModule);
        m_shaderModule = nullptr;
    }
}

bool ComputePipeline::rebuild(size_t codesize,
                              const void * const codedata)
{
    if(codesize < 4 || !codedata){
        SPDLOG_ERROR(__FUNCTION__ " invalid input !");
        return false;
    }

    if(codesize % 4){
        SPDLOG_WARN(__FUNCTION__ " spirv code data may be invalid!");
    }

    //  shader module
    std::shared_ptr<vk::ShaderModule> shadermodule(nullptr);
    if(!createShaderModule(codesize, codedata, shadermodule)){
        SPDLOG_ERROR("Create shader module failed!");
        return false;
    }

    //  pipelinelayout
    std::shared_ptr<vk::PipelineLayout> pipelinelayout(nullptr);
    if(!createPipelinelayout(m_descrSetLayouts,
                             m_pushConsts,
                             pipelinelayout
                             )){
        SPDLOG_ERROR("Create shader module failed!");
        return false;
    }


    //  pipeline cache
    auto pipelinecache = std::shared_ptr<vk::PipelineCache>();
    if(!m_useGlobalPipelinecache){
        pipelinecache = std::make_shared<vk::PipelineCache>();
        vk::PipelineCacheCreateInfo pipelineCacheInfo =
                vk::PipelineCacheCreateInfo();
        auto result = m_device->createPipelineCache(
                    &pipelineCacheInfo,
                    nullptr,
                    pipelinecache.get()
                    );
        if(result != vk::Result::eSuccess){
            SPDLOG_ERROR("Create pipelinecache failed!");
        }
    }

    //  pipeline
    auto pipeline = createPipeline(
                *pipelinelayout,
                *shadermodule,
                m_specConsts,
                pipelinecache
                );
    if(!pipeline){
        SPDLOG_ERROR("Create pipeline failed!");
        return false;
    }

    //  保存一些中间对象到成员变量中

    m_spirvCode.clear();
    m_spirvCode.assign(
                reinterpret_cast<const uint8_t * const>(codedata),
                reinterpret_cast<const uint8_t * const>(codedata) + codesize);

    if(m_shaderModule){
        m_device->destroyShaderModule(*m_shaderModule);
        m_shaderModule = nullptr;
    }
    m_shaderModule = shadermodule;

    if(m_layout){
        m_device->destroyPipelineLayout(*m_layout);
        m_layout = nullptr;
    }
    m_layout = pipelinelayout;

    if(m_pipelineCache){
        m_device->destroyPipelineCache(*m_pipelineCache);
        m_pipelineCache = nullptr;
    }
    m_pipelineCache = pipelinecache;

    if(m_pipeline){
        m_device->destroyPipeline(*m_pipeline);
        m_pipeline = nullptr;
    }
    m_pipeline = pipeline;

    return true;
}

bool ComputePipeline::rebuild()
{
    if(m_spirvCode.size() < 4){
        SPDLOG_ERROR("No valid spirv code data!");
        return false;
    }
    return rebuild(m_spirvCode.size(),
                   m_spirvCode.data());
}

bool ComputePipeline::createShaderModule(size_t codesize,
                                         const void * const codedata,
                                         std::shared_ptr<vk::ShaderModule> &shadermodule)
{
    vk::ShaderModuleCreateInfo shaderModuleInfo(
                vk::ShaderModuleCreateFlags(),
                codesize,
                reinterpret_cast<const uint32_t* const>(codedata)
                );

    shadermodule = std::make_shared<vk::ShaderModule>();
    auto result = m_device->createShaderModule(
                &shaderModuleInfo, nullptr, shadermodule.get());
    if(result != vk::Result::eSuccess){
        return false;
    }
    return true;
}

bool ComputePipeline::createPipelinelayout(
        std::vector<std::shared_ptr<vk::DescriptorSetLayout>> descrLayouts,
        std::shared_ptr<PipelineConstSZ> pushConsts,
        std::shared_ptr<vk::PipelineLayout> & pipelinelayout
        )
{
    vk::PipelineLayoutCreateInfo pipelineLayoutCreateInfo;
    std::vector<vk::DescriptorSetLayout> descrsetlayouts;
    descrsetlayouts.resize(descrLayouts.size());
    if(descrLayouts.size() > 0){
        std::transform(descrLayouts.begin(),
                       descrLayouts.end(),
                       descrsetlayouts.begin(),
                       [](std::shared_ptr<vk::DescriptorSetLayout> item){
            return *item;
        }
        );
        pipelineLayoutCreateInfo.setPSetLayouts(
                    descrsetlayouts.data()
                    );
        pipelineLayoutCreateInfo.setSetLayoutCount(
                    static_cast<uint32_t>(descrsetlayouts.size())
                    );
    }else{
        pipelineLayoutCreateInfo.setPSetLayouts(nullptr);
        pipelineLayoutCreateInfo.setSetLayoutCount(0);
    }

    pipelineLayoutCreateInfo.setFlags(vk::PipelineLayoutCreateFlags());

    vk::PushConstantRange pushConstantRange;
    pushConstantRange.setStageFlags(vk::ShaderStageFlagBits::eCompute);
    if(pushConsts){
        pushConsts->getPushConstantRange(pushConstantRange);
        pipelineLayoutCreateInfo.setPushConstantRangeCount(1);
        pipelineLayoutCreateInfo.setPPushConstantRanges(&pushConstantRange);
    }

    pipelinelayout = std::make_shared<vk::PipelineLayout>();
    auto result = m_device->createPipelineLayout(
                &pipelineLayoutCreateInfo,
                nullptr,
                pipelinelayout.get()
                );
    if(result != vk::Result::eSuccess){

        return false;
    }

    return true;
}

std::shared_ptr<vk::Pipeline> ComputePipeline::createPipeline(
        vk::PipelineLayout &pipelineLayout,
        vk::ShaderModule &shaderModule,
        std::shared_ptr<PipelineConstSZ> specConsts,
        std::shared_ptr<vk::PipelineCache> & pipelineCache)
{
    auto pipeline = std::make_shared<vk::Pipeline>();
    vk::ComputePipelineCreateInfo pipelineInfo;
    pipelineInfo.setFlags(vk::PipelineCreateFlags());
    pipelineInfo.setLayout(pipelineLayout);
    pipelineInfo.setBasePipelineHandle(vk::Pipeline()); // 可实现管道继承

    vk::PipelineShaderStageCreateInfo shaderStage(
                vk::PipelineShaderStageCreateFlags(),
                vk::ShaderStageFlagBits::eCompute,
                shaderModule,
                "main",
                nullptr);
    std::vector<vk::SpecializationMapEntry> specializationEntries;
    vk::SpecializationInfo specializationInfo(
                0,
                nullptr,
                0,
                nullptr);
    if(specConsts){
        specConsts->getSpecConstEntries(specializationEntries);
        if(specializationEntries.size() > 0){
            specializationInfo.setMapEntryCount(
                        static_cast<uint32_t>(specializationEntries.size()));
            specializationInfo.setPMapEntries(
                        specializationEntries.data());
            specializationInfo.setDataSize(
                        specConsts->size());
            specializationInfo.setPData(
                        specConsts->data()
                        );
            shaderStage.setPSpecializationInfo(
                        &specializationInfo);
        }
    }
    pipelineInfo.setStage(shaderStage);

    vk::PipelineCache * ppipelinecache = nullptr;
    if(m_useGlobalPipelinecache){
        ppipelinecache = m_GPipelineCache.get();
    }else if(pipelineCache){
        ppipelinecache = pipelineCache.get();
    }

    //    if(!ppipelinecache){
    //        pipelineCache = std::make_shared<vk::PipelineCache>();
    //        ppipelinecache = pipelineCache.get();
    //    }

    //        {
    //            std::vector<std::uint8_t> cached =
    //                    m_device->getPipelineCacheData(*ppipelinecache);
    //            SPDLOG_INFO("PipelineCache size={}", cached.size());
    //        }

    auto result = m_device->createComputePipeline(
                *ppipelinecache,
                pipelineInfo
                );
    if(result.result != vk::Result::eSuccess){
        return nullptr;
    }

    //        {
    //            std::vector<std::uint8_t> cached =
    //                    m_device->getPipelineCacheData(*ppipelinecache);
    //            SPDLOG_INFO("PipelineCache size={}", cached.size());
    //        }

    *pipeline = result.value;
    return pipeline;
}
